import os
import ahocorasick


class QuestClassifier:
    def __init__(self):
        cur_dir = '/'.join(os.path.abspath(__file__).split('/')[:-1])
        # 　特征词路径
        self.goods_path = os.path.join(cur_dir, 'dict/goods.txt')
        self.crowd_path = os.path.join(cur_dir, 'dict/crowd.txt')
        self.brand_path = os.path.join(cur_dir, 'dict/brand.txt')
        self.category_path = os.path.join(cur_dir, 'dict/category.txt')
        self.theme_path = os.path.join(cur_dir, 'dict/theme.txt')
        self.scene_path = os.path.join(cur_dir, 'dict/scene.txt')
        # 加载特征词
        self.goods_wds = [i.strip() for i in open(self.goods_path, encoding='utf8') if i.strip()]
        self.crowd_wds = [i.strip() for i in open(self.crowd_path, encoding='utf8') if i.strip()]
        self.brand_wds = [i.strip() for i in open(self.brand_path, encoding='utf8') if i.strip()]
        self.category_wds = [i.strip() for i in open(self.category_path, encoding='utf8') if i.strip()]
        self.theme_wds = [i.strip() for i in open(self.theme_path, encoding='utf8') if i.strip()]
        self.scene_wds = [i.strip() for i in open(self.scene_path, encoding='utf8') if i.strip()]
        self.region_words = set(self.goods_wds + self.brand_wds + self.category_wds + self.theme_wds + self.scene_wds + self.crowd_wds)
        # 构造领域actree
        self.region_tree = self.build_actree(list(self.region_words))
        # 构建词典
        self.wdtype_dict = self.build_wdtype_dict()

        self.brand_pwds = ['品牌', '牌子']
        self.crowd_pwds = ['适用', '针对', '合适', '适合', '面向']

        print('model init finished ......')
        return

    '''分类主函数'''

    def classify(self, question):
        data = {}
        entity_dict = self.check_entity(question)
        if not entity_dict:
            return {}
        data['args'] = entity_dict
        # 收集问句当中所涉及到的实体类型
        types = []
        for type_ in entity_dict.values():
            types += type_
        # print(types)
        question_types = []

        # 匹配问题类型
        # 商品品牌
        if self.check_words(self.brand_pwds, question) and ('goods' in types):
            question_type = 'goods_brand'
            question_types.append(question_type)
        # 已知品牌找商品
        elif 'brand' in types:
            question_type = 'brand_goods'
            question_types.append(question_type)
        # 根据人群找商品
        elif self.check_words(self.crowd_pwds, question) and ('crowd' in types):
            question_type = 'crowd_goods'
            question_types.append(question_type)
        # 根据分类找商品
        elif 'category' in types:
            question_type = 'category_goods'
            question_types.append(question_type)
        # 根据主题找商品
        elif 'theme' in types:
            question_type = 'theme_goods'
            question_types.append(question_type)
        # 根据场景找商品
        elif 'scene' in types:
            question_type = 'scene_goods'
            question_types.append(question_type)

        # 将多个分类结果进行合并处理，组装成一个字典
        data['question_types'] = question_types

        return data

    '''构造词对应的类型'''

    def build_wdtype_dict(self):
        wd_dict = dict()
        for wd in self.region_words:
            wd_dict[wd] = []
            if wd in self.goods_wds:
                wd_dict[wd].append('goods')
            if wd in self.brand_wds:
                wd_dict[wd].append('brand')
            if wd in self.crowd_wds:
                wd_dict[wd].append('crowd')
            if wd in self.category_wds:
                wd_dict[wd].append('category')
            if wd in self.theme_wds:
                wd_dict[wd].append('theme')
            if wd in self.scene_wds:
                wd_dict[wd].append('scene')
        return wd_dict

    '''问句过滤'''

    def check_entity(self, question):
        region_wds = []
        for i in self.region_tree.iter(question):
            wd = i[1][1]
            region_wds.append(wd)
        stop_wds = []
        for wd1 in region_wds:
            for wd2 in region_wds:
                if wd1 in wd2 and wd1 != wd2:
                    stop_wds.append(wd1)
        final_wds = [i for i in region_wds if i not in stop_wds]
        final_dict = {i: self.wdtype_dict.get(i) for i in final_wds}

        return final_dict

    '''构造actree，加速过滤'''

    def build_actree(self, wordlist):
        actree = ahocorasick.Automaton()
        for index, word in enumerate(wordlist):
            actree.add_word(word, (index, word))
        actree.make_automaton()
        return actree

    '''基于特征词进行分类'''

    def check_words(self, wds, sent):
        for wd in wds:
            if wd in sent:
                return True
        return False


if __name__ == '__main__':
    handler = QuestClassifier()
    while 1:
        question = input('input an question:')
        data = handler.classify(question)
        print(data)
